#include "tools.h"
#include "faiss_index.h"
int main(int argc, char* argv[])
{
	caffe::GlobalInit(&argc,&argv);
	if (argc < 3) {
	//	printHelp();
		return -1;
	}
	if (!strcmp(argv[1], "test")) 
	{
		// application.exe feature SampleImages.lst SampleImages.feature
		time_t t1 = time(NULL);

		cout<<Indexsingle::initance()<<endl;
		Indexsingle::update("sample.lst","sample.feature");
		//Indexsingle::initance()->appendFeature("sample.lst","sample.feature");
		//Indexsingle::initance()->save();

		time_t t2 = time(NULL);
		cout << "extract feature done, time: " << (t2 - t1) << " seconds." << endl;
	}
	if (!strcmp(argv[1], "feature")) 
	{
		// application.exe feature SampleImages.lst SampleImages.feature
		time_t t1 = time(NULL);

		int device_id = 0;
		extract_feat_saved(argv[2], argv[3], Mode, device_id);

		time_t t2 = time(NULL);
	//	cout << "extract feature done, time: " << (t2 - t1) << " seconds." << endl;
	}

	if(!strcmp(argv[1], "append"))
	{
		// application.exe append append.lst retrival.feature retrival.lst
		ifstream ifs(argv[2]); // append list
		ofstream ofs(argv[4], ios::app); // retrival list
		string line;
		while (std::getline(ifs, line))
		{
			ofs << line << endl;
		}
		ifs.close();
		ofs.close();

		int device_id = 0;
		bool append_feature = true;
		extract_feat_saved(argv[2], argv[3], Mode, device_id, append_feature);
	}

	if(!strcmp(argv[1], "remove"))
        {
		// application remove remove.lst retrival.feature retrival.lst
		ifstream ifs1(argv[2]); // remove list
		ifstream ifs2(argv[4]); // retrival list

		string line;
		vector<int> remove_pos;
		vector<string> retr_lst;
		vector<string> retr_bak;
		while (std::getline(ifs2, line)) // read retrival list
		{
			retr_lst.push_back(line);
			retr_bak.push_back(line);
		}
                ifs2.close();
		
		while (std::getline(ifs1, line)) // read remove list
                {
			vector<string>::iterator iter = find(retr_lst.begin(), retr_lst.end(), line);
			if (iter != retr_lst.end())
			{
				int pos = distance(retr_lst.begin(), iter);
				remove_pos.push_back(pos);

				vector<string>::iterator iter1 = find(retr_bak.begin(), retr_bak.end(), line);
				if (iter1 != retr_bak.end())	
					retr_bak.erase(iter1);
			}
                }
                ifs1.close();
	
		ofstream ofs(argv[4]); // retrival list after removed
		for (int i = 0; i < retr_bak.size(); i++)
		{
			line = retr_bak[i];
			ofs << line << endl;
		}	
		ofs.close();

		/***************************************/

		FILE *in1 = fopen(argv[3], "rb");
                int retrivalNum, featDim;
                fread(&retrivalNum, sizeof(int), 1, in1);
                fread(&featDim, sizeof(int), 1, in1);
                cout<<"retrival num: "<<retrivalNum<<" feature dim: "<<featDim<<endl;

		retrivalNum = retr_lst.size();  

		FILE *out1 = fopen("temp.feature", "wb");
		int newNum = retr_bak.size();
		fwrite(&newNum, sizeof(int), 1, out1);
    		fwrite(&featDim, sizeof(int), 1, out1);

		int index = 0;
		sort(remove_pos.begin(), remove_pos.end());
		for(int i = 0; i < retrivalNum; i++)
                {
                        float *temp = new float[featDim];
                        fread(temp, sizeof(float), featDim, in1);
               		
			if(i == remove_pos[index] && index < remove_pos.size())
				index++;
			else fwrite(temp, sizeof(float), featDim, out1);
		 
                }
                fclose(in1);
                fclose(out1);
                
                if( rename("temp.feature", argv[3]) == 0) 
                        cout << "success" << endl;
                else    cout << "fail" << endl;

	}

	if (!strcmp(argv[1], "batchquery")) 
	{
	//	cout<<argv[0]<<" "<<argv[1]<<" "<<argv[2]<<" "<<argv[3]<<" "<<argv[4]<<endl;
		// application.exe batchquery QueryImage.lst Images.feature Image.lst 
		if(argc != 5)
   		{
			//cout << "Parameter error!" << endl << "application.exe batchquery [QueryImage.lst]";
			//cout << " [RetrivalImages.feature] [RetrivalImage.lst] [text.feature] [text.lst]" << endl;
        		return -1;
   		}
		// read image list
		time_t t1 = time(NULL);	
		ifstream ifs1(argv[2]);
                ifstream ifs2(argv[4]);
                string line;
                vector<string> query_name;
                vector<string> retrival_name;
                while (std::getline(ifs1, line))
                       query_name.push_back(line_to_line(line));
                ifs1.close();
                while (std::getline(ifs2, line))
                       retrival_name.push_back(line_after_line(line));
                ifs2.close();

		
                //cout<<"query num: "<<query_name.size()<<" retrival num: "<<retrival_name.size()<<endl;
		//cout<<"text image num: "<<text_name.size()<<endl;
/**********************************/
		const int feat_start_pos = 2048; //###
		const int feat_end_pos = 2560;   //###
		const int feat_width = feat_end_pos - feat_start_pos; //###

	
		// extract query image feature
		int device_id = 0;
		vector<bool> is_image;
		Mat feat_query = extract_feat_not_saved(argv[2],is_image, Mode, device_id);
		feat_query = feat_query.colRange(feat_start_pos, feat_end_pos).clone(); //###
		int queryNum = feat_query.rows;
		int featDim = feat_query.cols; 
		//cout<<"query num: "<<queryNum<<" feature dim: "<<featDim<<endl;

		for(int i = 0; i < queryNum; i++)
		{
			Mat dataRow = feat_query.rowRange(i, i + 1).clone();
			float feat_norm = norm(dataRow);
			if(feat_norm != 0)
				dataRow = dataRow/feat_norm;
			dataRow.copyTo(feat_query.rowRange(i, i + 1));
		}	
		time_t t2 = time(NULL);		
		//cout << "extract query feature done, time: " << (t2- t1) << " seconds."<< endl;

		FILE *in1 = fopen(argv[3], "rb");
		int retrivalNum;
                fread(&retrivalNum, sizeof(int), 1, in1);
                fread(&featDim, sizeof(int), 1, in1);
                //cout<<"retrival num: "<<retrivalNum<<" feature dim: "<<featDim<<endl;

		retrivalNum = retrival_name.size();
                Mat feat_retrival = Mat::zeros(retrivalNum, feat_width, CV_32FC1); //###
                for(int i = 0; i < retrivalNum; i++)
                {
                        float *temp = new float[featDim];
                        fread(temp, sizeof(float), featDim, in1);

                        Mat dataRow(1, featDim, CV_32FC1, temp);
			dataRow = dataRow.colRange(feat_start_pos, feat_end_pos).clone(); //###
                        float feat_norm = norm(dataRow);
                        if(feat_norm != 0)
                                dataRow = dataRow/feat_norm;
                        dataRow.copyTo(feat_retrival.rowRange(i, i + 1));
                }
                fclose(in1);
		time_t t3 = time(NULL);
		//cout << "text feature read done, time: " << (t3- t2) << " seconds."<< endl;


/***********************************************************************/

		time_t tupian = time(NULL);
		cosine_retrival(is_image,query_name, retrival_name, queryNum, retrivalNum, feat_query, feat_retrival, 0.845);
		time_t end = time(NULL);


		cout << "total time: " << end - t1 << " seconds." << endl;	
/**************************************************************************/
	}

	if (!strcmp(argv[1], "faiss")) 
	{
	//	cout<<argv[0]<<" "<<argv[1]<<" "<<argv[2]<<" "<<argv[3]<<" "<<argv[4]<<endl;
		// application.exe batchquery QueryImage.lst Images.feature Image.lst 
		if(argc != 5)
   		{
			//cout << "Parameter error!" << endl << "application.exe batchquery [QueryImage.lst]";
			//cout << " [RetrivalImages.feature] [RetrivalImage.lst] [text.feature] [text.lst]" << endl;
        		return -1;
   		}
		// read image list
		time_t t1 = time(NULL);	
		ifstream ifs1(argv[2]);
                ifstream ifs2(argv[4]);
                string line;
                vector<string> query_name;
                vector<string> retrival_name;
                while (std::getline(ifs1, line))
                       query_name.push_back(line_to_line(line));
                ifs1.close();
                while (std::getline(ifs2, line))
                       retrival_name.push_back(line_after_line(line));
                ifs2.close();

		
                //cout<<"query num: "<<query_name.size()<<" retrival num: "<<retrival_name.size()<<endl;
		//cout<<"text image num: "<<text_name.size()<<endl;
/**********************************/
		const int feat_start_pos = 2048; //###
		const int feat_end_pos = 2560;   //###
		const int feat_width = feat_end_pos - feat_start_pos; //###

	
		// extract query image feature
		int device_id = 0;
		vector<bool> is_image;
		Mat feat_query = extract_feat_not_saved(argv[2], is_image,Mode, device_id);
		feat_query = feat_query.colRange(feat_start_pos, feat_end_pos).clone(); //###
		int queryNum = feat_query.rows;
		int featDim = feat_query.cols; 
		//cout<<"query num: "<<queryNum<<" feature dim: "<<featDim<<endl;

		for(int i = 0; i < queryNum; i++)
		{
			Mat dataRow = feat_query.rowRange(i, i + 1).clone();
			float feat_norm = norm(dataRow);
			if(feat_norm != 0)
				dataRow = dataRow/feat_norm;
			dataRow.copyTo(feat_query.rowRange(i, i + 1));
		}	
		time_t t2 = time(NULL);		
		//cout << "extract query feature done, time: " << (t2- t1) << " seconds."<< endl;

		FILE *in1 = fopen(argv[3], "rb");
		int retrivalNum;
                fread(&retrivalNum, sizeof(int), 1, in1);
                fread(&featDim, sizeof(int), 1, in1);
                //cout<<"retrival num: "<<retrivalNum<<" feature dim: "<<featDim<<endl;

		retrivalNum = retrival_name.size();
                Mat feat_retrival = Mat::zeros(retrivalNum, feat_width, CV_32FC1); //###
                for(int i = 0; i < retrivalNum; i++)
                {
                        float *temp = new float[featDim];
                        fread(temp, sizeof(float), featDim, in1);

                        Mat dataRow(1, featDim, CV_32FC1, temp);
			dataRow = dataRow.colRange(feat_start_pos, feat_end_pos).clone(); //###
                        float feat_norm = norm(dataRow);
                        if(feat_norm != 0)
                                dataRow = dataRow/feat_norm;
                        dataRow.copyTo(feat_retrival.rowRange(i, i + 1));
                }
                fclose(in1);
		//float *featArr=new float[feat_retrival.rows*feat_retrival.cols];
		float* featArr=(float*)feat_retrival.data;
                faiss::IndexFlatIP IPindex=faiss::IndexFlatIP(feat_width);
		IPindex.add(retrivalNum,featArr);
		time_t t3 = time(NULL);
		cout << "text feature read done, time: " << (t3- t2) << " seconds."<< endl;


/***********************************************************************/

		time_t tupian = time(NULL);
		cosine_retrival_faiss(IPindex,is_image,query_name, retrival_name, queryNum, retrivalNum, feat_query, feat_retrival, 0.845);
		time_t end = time(NULL);


		cout << "total time: " << end - t1 << " seconds." << endl;	
/**************************************************************************/
	}
	
	return 0;
}

